{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/athekunal/.local/lib/python3.10/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/home/athekunal/.local/lib/python3.10/site-packages/pyfolio/pos.py:26: UserWarning: Module \"zipline.assets\" not found; multipliers will not be applied to position notionals.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import ray\n",
    "assert ray.__version__ >='2.2.0', \"Please install ray 2.2.0 by doing 'pip install ray[rllib] ray[tune] lz4' , lz4 is for population based tuning\"\n",
    "\n",
    "#Importing the libraries\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# matplotlib.use('Agg')\n",
    "import ray.rllib.algorithms.ppo as ppo\n",
    "import ray.rllib.algorithms.a2c as a2c\n",
    "import ray.rllib.algorithms.a3c as a3c\n",
    "import ray.rllib.algorithms.td3 as td3\n",
    "import ray.rllib.algorithms.ddpg as ddpg\n",
    "import ray.rllib.algorithms.appo as appo\n",
    "import datetime\n",
    "%matplotlib inline\n",
    "from finrl import config\n",
    "from finrl.meta.preprocessor.yahoodownloader import YahooDownloader\n",
    "from finrl.meta.preprocessor.preprocessors import FeatureEngineer, data_split\n",
    "from finrl.meta.env_stock_trading.env_stocktrading_np import StockTradingEnv as StockTradingEnv_numpy \n",
    "from finrl.meta.data_processor import DataProcessor\n",
    "from finrl.plot import backtest_stats, backtest_plot, get_daily_return, get_baseline\n",
    "import ray\n",
    "from pprint import pprint\n",
    "# from ray.rllib.algorithms.ppo import ppo\n",
    "# from ray.rllib.algorithms.ddpg import ddpg\n",
    "# from ray.rllib.algorithms.a2c import a2c\n",
    "# from ray.rllib.algorithms.ddpg import ddpg,td3\n",
    "# from ray.rllib.algorithms import ddpg\n",
    "import ray.rllib.algorithms.ppo as ppo\n",
    "# from ray.rllib.algorithms.sac import sac\n",
    "import sys\n",
    "sys.path.append(\"../FinRL-Library\")\n",
    "import os\n",
    "import itertools\n",
    "from ray import tune\n",
    "from ray.tune.search import ConcurrencyLimiter\n",
    "from ray.tune.schedulers import AsyncHyperBandScheduler\n",
    "from ray.tune.search.optuna import OptunaSearch\n",
    "from ray.tune.schedulers import ASHAScheduler\n",
    "from ray.tune.registry import register_env\n",
    "from ray import air\n",
    "from ray.air import session\n",
    "import time\n",
    "import psutil\n",
    "psutil_memory_in_bytes = psutil.virtual_memory().total\n",
    "ray._private.utils.get_system_memory = lambda: psutil_memory_in_bytes\n",
    "from typing import Dict, Optional, Any\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @Author: Astarag Mohapatra\n",
    "\n",
    "import ray\n",
    "\n",
    "assert (\n",
    "    ray.__version__ > \"2.0.0\"\n",
    "), \"Please install ray 2.2.0 by doing 'pip install ray[rllib] ray[tune] lz4' , lz4 is for population based tuning\"\n",
    "from pprint import pprint\n",
    "\n",
    "from ray import tune\n",
    "from ray.tune.search import ConcurrencyLimiter\n",
    "from ray.rllib.algorithms import Algorithm\n",
    "from ray.tune import register_env\n",
    "\n",
    "from ray.air import RunConfig, FailureConfig, ScalingConfig\n",
    "from ray.tune.tune_config import TuneConfig\n",
    "from ray.air.config import CheckpointConfig\n",
    "from ray.tune.callback import Callback\n",
    "\n",
    "import psutil\n",
    "\n",
    "psutil_memory_in_bytes = psutil.virtual_memory().total\n",
    "ray._private.utils.get_system_memory = lambda: psutil_memory_in_bytes\n",
    "from typing import Dict, Optional, Any, List, Union\n",
    "\n",
    "\n",
    "class DRLlibv2:\n",
    "    \"\"\"\n",
    "    It instantiates RLlib model with Ray tune functionality\n",
    "    Params\n",
    "    -------------------------------------\n",
    "    trainable:\n",
    "        Any Trainable class that takes config as parameter\n",
    "    train_env:\n",
    "        Training environment instance\n",
    "    train_env_name: str\n",
    "        Name of the training environment\n",
    "    params: dict\n",
    "        hyperparameters dictionary\n",
    "    run_name: str\n",
    "        tune run name\n",
    "    framework: str\n",
    "        \"torch\" or \"tf\" for tensorflow\n",
    "    local_dir: str\n",
    "         to save the results and tensorboard plots\n",
    "    num_workers: int\n",
    "        number of workers\n",
    "    search_alg\n",
    "        search space for hyperparameters\n",
    "    concurrent_trials:\n",
    "         Number of concurrent hyperparameters trial to run\n",
    "    num_samples: int\n",
    "         Number of samples of hyperparameters config to run\n",
    "    scheduler:\n",
    "        Stopping suboptimal trials\n",
    "    log_level: str = \"WARN\",\n",
    "        Verbosity: \"DEBUG\"\n",
    "    num_gpus: Union[float, int] = 0\n",
    "        GPUs for trial\n",
    "    num_cpus: Union[float, int] = 2\n",
    "        CPUs for rollout collection\n",
    "    dataframe_save: str\n",
    "        Saving the tune results\n",
    "    metric: str\n",
    "        Metric for hyperparameter optimization in Bayesian Methods\n",
    "    mode: str\n",
    "        Maximize or Minimize the metric\n",
    "    max_failures: int\n",
    "        Number of failures to TuneError\n",
    "    training_iterations: str\n",
    "         Number of times session.report() is called\n",
    "    checkpoint_num_to_keep: int\n",
    "        Number of checkpoints to keep\n",
    "    checkpoint_freq: int\n",
    "        Checkpoint freq wrt training iterations\n",
    "    reuse_actors:bool\n",
    "        Reuse actors for tuning\n",
    "    callbacks:\n",
    "        callbacks integration for ray tune\n",
    "\n",
    "    It has the following methods:\n",
    "    Methods\n",
    "    -------------------------------------\n",
    "        train_tune_model: It takes in the params dictionary and fits in sklearn style to our trainable class\n",
    "        restore_agent: It restores previously errored or stopped trials or experiments\n",
    "        infer_results: It returns the results dataframe and trial informations\n",
    "        get_test_agent: It returns the testing agent for inference\n",
    "\n",
    "    Example\n",
    "    ---------------------------------------\n",
    "    def sample_ppo_params():\n",
    "        return {\n",
    "            \"entropy_coeff\": tune.loguniform(0.00000001, 0.1),\n",
    "            \"lr\": tune.loguniform(5e-5, 0.001),\n",
    "            \"sgd_minibatch_size\": tune.choice([ 32, 64, 128, 256, 512]),\n",
    "            \"lambda\": tune.choice([0.1,0.3,0.5,0.7,0.9,1.0]),\n",
    "        }\n",
    "        #Tree Parzen Estimator\n",
    "    optuna_search = OptunaSearch(\n",
    "        metric=\"episode_reward_mean\",\n",
    "        mode=\"max\")\n",
    "    drl_agent = DRLlibv2(\n",
    "        trainable=\"PPO\",\n",
    "        train_env=env(train_env_config),\n",
    "        train_env_name=\"StockTrading_train\",\n",
    "        framework=\"torch\",\n",
    "        num_workers=1,\n",
    "        log_level=\"DEBUG\",\n",
    "        run_name = 'test',\n",
    "        local_dir = \"test\",\n",
    "        params = sample_ppo_params(),\n",
    "        num_samples = 1,\n",
    "        num_gpus=1,\n",
    "        training_iterations=10,\n",
    "        search_alg = optuna_search,\n",
    "        checkpoint_freq=5\n",
    "    )\n",
    "    #Tune or train the model\n",
    "    res = drl_agent.train_tune_model()\n",
    "\n",
    "    #Get the tune results\n",
    "    results_df, best_result = drl_agent.infer_results()\n",
    "\n",
    "    #Get the best testing agent\n",
    "    test_agent = drl_agent.get_test_agent(test_env_instance,'StockTrading_testenv')\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        trainable: Union[str, Any],\n",
    "        params: dict,\n",
    "        train_env=None,\n",
    "        train_env_name: str='',\n",
    "        run_name: str = \"tune_run\",\n",
    "        framework: str = \"torch\",\n",
    "        local_dir: str = \"tune_results\",\n",
    "        num_workers: int = 1,\n",
    "        search_alg=None,\n",
    "        concurrent_trials: int = 0,\n",
    "        num_samples: int = 0,\n",
    "        scheduler=None,\n",
    "        log_level: str = \"WARN\",\n",
    "        num_gpus: Union[float, int] = 0,\n",
    "        num_cpus: Union[float, int] = 2,\n",
    "        dataframe_save: str = \"tune.csv\",\n",
    "        metric: str = \"episode_reward_mean\",\n",
    "        mode: Union[str, List[str]] = \"max\",\n",
    "        max_failures: int = 0,\n",
    "        training_iterations: int = 100,\n",
    "        checkpoint_num_to_keep: Union[None, int] = None,\n",
    "        checkpoint_freq: int = 0,\n",
    "        reuse_actors: bool = False,\n",
    "        callbacks:Optional[List[\"Callback\"]]=None\n",
    "    ):\n",
    "\n",
    "        if train_env is not None:register_env(train_env_name, lambda config: train_env)\n",
    "\n",
    "        self.params = params\n",
    "        self.params[\"framework\"] = framework\n",
    "        self.params[\"log_level\"] = log_level\n",
    "        self.params[\"num_gpus\"] = num_gpus\n",
    "        self.params[\"num_workers\"] = num_workers\n",
    "        self.params[\"env\"] = train_env_name\n",
    "\n",
    "        self.run_name = run_name\n",
    "        self.local_dir = local_dir\n",
    "        self.search_alg = search_alg\n",
    "        if concurrent_trials != 0:\n",
    "            self.search_alg = ConcurrencyLimiter(\n",
    "                self.search_alg, max_concurrent=concurrent_trials\n",
    "            )\n",
    "        self.scheduler = scheduler\n",
    "        self.num_samples = num_samples\n",
    "        self.trainable = trainable\n",
    "        # self.trainable = tune.with_resources(self.trainable(),{\"cpu\":num_cpus,\"gpu\":num_gpus})\n",
    "        if isinstance(self.trainable, str):\n",
    "            self.trainable.upper()\n",
    "        self.num_cpus = num_cpus\n",
    "        self.num_gpus = num_gpus\n",
    "        self.dataframe_save = dataframe_save\n",
    "        self.metric = metric\n",
    "        self.mode = mode\n",
    "        self.max_failures = max_failures\n",
    "        self.training_iterations = training_iterations\n",
    "        self.checkpoint_freq = checkpoint_freq\n",
    "        self.checkpoint_num_to_keep = checkpoint_num_to_keep\n",
    "        self.reuse_actors = reuse_actors\n",
    "        self.callbacks = callbacks\n",
    "\n",
    "    def train_tune_model(self):\n",
    "        \"\"\"\n",
    "        Tuning and training the model\n",
    "        Returns the results object\n",
    "        \"\"\"\n",
    "        ray.init(\n",
    "            num_cpus=self.num_cpus, num_gpus=self.num_gpus, ignore_reinit_error=True\n",
    "        )\n",
    "\n",
    "        tuner = tune.Tuner(\n",
    "            self.trainable,\n",
    "            param_space=self.params,\n",
    "            tune_config=TuneConfig(\n",
    "                search_alg=self.search_alg,\n",
    "                num_samples=self.num_samples,\n",
    "                metric=self.metric,\n",
    "                mode=self.mode,\n",
    "                reuse_actors=self.reuse_actors,\n",
    "            ),\n",
    "            run_config=RunConfig(\n",
    "                name=self.run_name,\n",
    "                local_dir=self.local_dir,\n",
    "                callbacks=self.callbacks,\n",
    "                failure_config=FailureConfig(\n",
    "                    max_failures=self.max_failures, fail_fast=False\n",
    "                ),\n",
    "                stop={\"training_iteration\": self.training_iterations},\n",
    "                checkpoint_config=CheckpointConfig(\n",
    "                    num_to_keep=self.checkpoint_num_to_keep,\n",
    "                    checkpoint_score_attribute=self.metric,\n",
    "                    checkpoint_score_order=self.mode,\n",
    "                    checkpoint_frequency=self.checkpoint_freq,\n",
    "                    checkpoint_at_end=True,\n",
    "                ),\n",
    "                verbose=3,\n",
    "            ),\n",
    "        )\n",
    "\n",
    "        self.results = tuner.fit()\n",
    "        # if self.search_alg is not None: self.search_alg.save_to_dir(self.local_dir)\n",
    "        # ray.shutdown()\n",
    "        return self.results\n",
    "\n",
    "    def infer_results(self, to_dataframe: str = None, mode: str = \"a\"):\n",
    "        \"\"\"\n",
    "        Get tune results in a dataframe and best results object\n",
    "        \"\"\"\n",
    "        results_df = self.results.get_dataframe()\n",
    "\n",
    "        if to_dataframe is None:\n",
    "            to_dataframe = self.dataframe_save\n",
    "\n",
    "        results_df.to_csv(to_dataframe, mode=mode)\n",
    "\n",
    "        best_result = self.results.get_best_result()\n",
    "        # best_result = self.results.get_best_result()\n",
    "        # best_metric = best_result.metrics\n",
    "        # best_checkpoint = best_result.checkpoint\n",
    "        # best_trial_dir = best_result.log_dir\n",
    "        # results_df = self.results.get_dataframe()\n",
    "\n",
    "        return results_df, best_result\n",
    "\n",
    "    def restore_agent(\n",
    "        self,\n",
    "        checkpoint_path: str = \"\",\n",
    "        restore_search: bool = False,\n",
    "        resume_unfinished: bool = True,\n",
    "        resume_errored: bool = False,\n",
    "        restart_errored: bool = False,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Restore errored or stopped trials\n",
    "        \"\"\"\n",
    "        # if restore_search:\n",
    "        # self.search_alg = self.search_alg.restore_from_dir(self.local_dir)\n",
    "        if checkpoint_path == \"\":\n",
    "            checkpoint_path = self.results.get_best_result().checkpoint._local_path\n",
    "\n",
    "        restored_agent = tune.Tuner.restore(\n",
    "            checkpoint_path,\n",
    "            restart_errored=restart_errored,\n",
    "            resume_unfinished=resume_unfinished,\n",
    "            resume_errored=resume_errored,\n",
    "        )\n",
    "        print(restored_agent)\n",
    "        self.results = restored_agent.fit()\n",
    "\n",
    "        # self.search_alg.save_to_dir(self.local_dir)\n",
    "        return self.results\n",
    "\n",
    "    def get_test_agent(self, test_env, test_env_name: str, checkpoint=None):\n",
    "        \"\"\"\n",
    "        Get test agent\n",
    "        \"\"\"\n",
    "        if test_env is not None:register_env(test_env_name, lambda config: test_env)\n",
    "\n",
    "        if checkpoint is None:\n",
    "            checkpoint = self.results.get_best_result().checkpoint\n",
    "\n",
    "        testing_agent = Algorithm.from_checkpoint(checkpoint)\n",
    "        # testing_agent.config['env'] = test_env_name\n",
    "\n",
    "        return testing_agent\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BUILDING THE ENVIRONMENT\n",
    "\n",
    "* Let's build an environment to test our algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_train_env(start_date, end_date, ticker_list, data_source, time_interval, \n",
    "          technical_indicator_list, env, model_name, if_vix = True,\n",
    "          **kwargs):\n",
    "    \n",
    "    #fetch data\n",
    "    DP = DataProcessor(data_source, **kwargs)\n",
    "    data = DP.download_data(ticker_list, start_date, end_date, time_interval)\n",
    "    data = DP.clean_data(data)\n",
    "    data = DP.add_technical_indicator(data, technical_indicator_list)\n",
    "    if if_vix:\n",
    "        data = DP.add_vix(data)\n",
    "    price_array, tech_array, turbulence_array = DP.df_to_array(data, if_vix)\n",
    "    train_env_config = {'price_array':price_array,\n",
    "              'tech_array':tech_array,\n",
    "              'turbulence_array':turbulence_array,\n",
    "              'if_train':True}\n",
    "    \n",
    "    return train_env_config\n",
    "\n",
    "def calculate_sharpe(episode_reward:list):\n",
    "  perf_data = pd.DataFrame(data=episode_reward,columns=['reward'])\n",
    "  perf_data['daily_return'] = perf_data['reward'].pct_change(1)\n",
    "  if perf_data['daily_return'].std() !=0:\n",
    "    sharpe = (252**0.5)*perf_data['daily_return'].mean()/ \\\n",
    "          perf_data['daily_return'].std()\n",
    "    return sharpe\n",
    "  else:\n",
    "    return 0\n",
    "\n",
    "def get_test_config(start_date, end_date, ticker_list, data_source, time_interval, \n",
    "         technical_indicator_list, env, model_name, if_vix = True,\n",
    "         **kwargs):\n",
    "  \n",
    "  DP = DataProcessor(data_source, **kwargs)\n",
    "  data = DP.download_data(ticker_list, start_date, end_date, time_interval)\n",
    "  data = DP.clean_data(data)\n",
    "  data = DP.add_technical_indicator(data, technical_indicator_list)\n",
    "  \n",
    "  if if_vix:\n",
    "      data = DP.add_vix(data)\n",
    "  \n",
    "  price_array, tech_array, turbulence_array = DP.df_to_array(data, if_vix)\n",
    "  test_env_config = {'price_array':price_array,\n",
    "            'tech_array':tech_array,\n",
    "            'turbulence_array':turbulence_array,'if_train':False}\n",
    "  return test_env_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_START_DATE = '2014-01-01'\n",
    "TRAIN_END_DATE = '2019-07-30'\n",
    "\n",
    "VAL_START_DATE = '2019-08-01'\n",
    "VAL_END_DATE = '2021-07-30'\n",
    "\n",
    "TEST_START_DATE = '2021-08-01'\n",
    "TEST_END_DATE = '2023-02-01'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from finrl.config_tickers import DOW_30_TICKER\n",
    "technical_indicator_list =config.INDICATORS\n",
    "\n",
    "model_name = 'PPO'\n",
    "env = StockTradingEnv_numpy\n",
    "ticker_list = DOW_30_TICKER\n",
    "data_source = 'yahoofinance'\n",
    "time_interval = '1D'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "Shape of DataFrame:  (55899, 9)\n",
      "Clean data for AAPL\n",
      "Data clean for AAPL is finished.\n",
      "Clean data for AMGN\n",
      "Data clean for AMGN is finished.\n",
      "Clean data for AXP\n",
      "Data clean for AXP is finished.\n",
      "Clean data for BA\n",
      "Data clean for BA is finished.\n",
      "Clean data for CAT\n",
      "Data clean for CAT is finished.\n",
      "Clean data for CRM\n",
      "Data clean for CRM is finished.\n",
      "Clean data for CSCO\n",
      "Data clean for CSCO is finished.\n",
      "Clean data for CVX\n",
      "Data clean for CVX is finished.\n",
      "Clean data for DIS\n",
      "Data clean for DIS is finished.\n",
      "Clean data for DOW\n",
      "NaN data on start date, fill using first valid data.\n",
      "Data clean for DOW is finished.\n",
      "Clean data for GS\n",
      "Data clean for GS is finished.\n",
      "Clean data for HD\n",
      "Data clean for HD is finished.\n",
      "Clean data for HON\n",
      "Data clean for HON is finished.\n",
      "Clean data for IBM\n",
      "Data clean for IBM is finished.\n",
      "Clean data for INTC\n",
      "Data clean for INTC is finished.\n",
      "Clean data for JNJ\n",
      "Data clean for JNJ is finished.\n",
      "Clean data for JPM\n",
      "Data clean for JPM is finished.\n",
      "Clean data for KO\n",
      "Data clean for KO is finished.\n",
      "Clean data for MCD\n",
      "Data clean for MCD is finished.\n",
      "Clean data for MMM\n",
      "Data clean for MMM is finished.\n",
      "Clean data for MRK\n",
      "Data clean for MRK is finished.\n",
      "Clean data for MSFT\n",
      "Data clean for MSFT is finished.\n",
      "Clean data for NKE\n",
      "Data clean for NKE is finished.\n",
      "Clean data for PG\n",
      "Data clean for PG is finished.\n",
      "Clean data for TRV\n",
      "Data clean for TRV is finished.\n",
      "Clean data for UNH\n",
      "Data clean for UNH is finished.\n",
      "Clean data for V\n",
      "Data clean for V is finished.\n",
      "Clean data for VZ\n",
      "Data clean for VZ is finished.\n",
      "Clean data for WBA\n",
      "Data clean for WBA is finished.\n",
      "Clean data for WMT\n",
      "Data clean for WMT is finished.\n",
      "Data clean all finished!\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "Shape of DataFrame:  (1907, 9)\n",
      "Clean data for ^VIX\n",
      "Data clean for ^VIX is finished.\n",
      "Data clean all finished!\n",
      "['AAPL' 'AMGN' 'AXP' 'BA' 'CAT' 'CRM' 'CSCO' 'CVX' 'DIS' 'DOW' 'GS' 'HD'\n",
      " 'HON' 'IBM' 'INTC' 'JNJ' 'JPM' 'KO' 'MCD' 'MMM' 'MRK' 'MSFT' 'NKE' 'PG'\n",
      " 'TRV' 'UNH' 'V' 'VZ' 'WBA' 'WMT']\n",
      "Successfully transformed into array\n"
     ]
    }
   ],
   "source": [
    "import warnings \n",
    "warnings.filterwarnings(\"ignore\", category=FutureWarning)\n",
    "\n",
    "train_env_config = get_train_env(TRAIN_START_DATE, VAL_END_DATE, \n",
    "                     ticker_list, data_source, time_interval, \n",
    "                        technical_indicator_list, env, model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from ray.tune.registry import register_env\n",
    "from ray.tune import register_env\n",
    "env_name = 'StockTrading_train_env'\n",
    "\n",
    "def reg_env(config):\n",
    "    return env(config)\n",
    "register_env(env_name, lambda : env(train_env_config))\n",
    "\n",
    "train_env_instance = env(train_env_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from drllibv2 import DRLlibv2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GETTING HYPERPARAMETERS"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The hyperparameters configuration can be generated from the Config file of all the algorithms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'extra_python_environs_for_driver': {},\n",
       " 'extra_python_environs_for_worker': {},\n",
       " 'num_gpus': 0,\n",
       " 'num_cpus_per_worker': 1,\n",
       " 'num_gpus_per_worker': 0,\n",
       " '_fake_gpus': False,\n",
       " 'custom_resources_per_worker': {},\n",
       " 'placement_strategy': 'PACK',\n",
       " 'eager_tracing': False,\n",
       " 'eager_max_retraces': 20,\n",
       " 'tf_session_args': {'intra_op_parallelism_threads': 2,\n",
       "  'inter_op_parallelism_threads': 2,\n",
       "  'gpu_options': {'allow_growth': True},\n",
       "  'log_device_placement': False,\n",
       "  'device_count': {'CPU': 1},\n",
       "  'allow_soft_placement': True},\n",
       " 'local_tf_session_args': {'intra_op_parallelism_threads': 8,\n",
       "  'inter_op_parallelism_threads': 8},\n",
       " 'env': None,\n",
       " 'env_config': {},\n",
       " 'observation_space': None,\n",
       " 'action_space': None,\n",
       " 'env_task_fn': None,\n",
       " 'render_env': False,\n",
       " 'clip_rewards': None,\n",
       " 'normalize_actions': True,\n",
       " 'clip_actions': False,\n",
       " 'disable_env_checking': False,\n",
       " 'num_envs_per_worker': 1,\n",
       " 'sample_collector': ray.rllib.evaluation.collectors.simple_list_collector.SimpleListCollector,\n",
       " 'sample_async': False,\n",
       " 'enable_connectors': False,\n",
       " 'rollout_fragment_length': 'auto',\n",
       " 'batch_mode': 'truncate_episodes',\n",
       " 'remote_worker_envs': False,\n",
       " 'remote_env_batch_wait_ms': 0,\n",
       " 'validate_workers_after_construction': True,\n",
       " 'ignore_worker_failures': False,\n",
       " 'recreate_failed_workers': False,\n",
       " 'restart_failed_sub_environments': False,\n",
       " 'num_consecutive_worker_failures_tolerance': 100,\n",
       " 'horizon': None,\n",
       " 'soft_horizon': False,\n",
       " 'no_done_at_end': False,\n",
       " 'preprocessor_pref': 'deepmind',\n",
       " 'observation_filter': 'NoFilter',\n",
       " 'synchronize_filters': True,\n",
       " 'compress_observations': False,\n",
       " 'enable_tf1_exec_eagerly': False,\n",
       " 'sampler_perf_stats_ema_coef': None,\n",
       " 'gamma': 0.99,\n",
       " 'lr': 5e-05,\n",
       " 'train_batch_size': 4000,\n",
       " 'model': {'_use_default_native_models': False,\n",
       "  '_disable_preprocessor_api': False,\n",
       "  '_disable_action_flattening': False,\n",
       "  'fcnet_hiddens': [256, 256],\n",
       "  'fcnet_activation': 'tanh',\n",
       "  'conv_filters': None,\n",
       "  'conv_activation': 'relu',\n",
       "  'post_fcnet_hiddens': [],\n",
       "  'post_fcnet_activation': 'relu',\n",
       "  'free_log_std': False,\n",
       "  'no_final_linear': False,\n",
       "  'vf_share_layers': False,\n",
       "  'use_lstm': False,\n",
       "  'max_seq_len': 20,\n",
       "  'lstm_cell_size': 256,\n",
       "  'lstm_use_prev_action': False,\n",
       "  'lstm_use_prev_reward': False,\n",
       "  '_time_major': False,\n",
       "  'use_attention': False,\n",
       "  'attention_num_transformer_units': 1,\n",
       "  'attention_dim': 64,\n",
       "  'attention_num_heads': 1,\n",
       "  'attention_head_dim': 32,\n",
       "  'attention_memory_inference': 50,\n",
       "  'attention_memory_training': 50,\n",
       "  'attention_position_wise_mlp_dim': 32,\n",
       "  'attention_init_gru_gate_bias': 2.0,\n",
       "  'attention_use_n_prev_actions': 0,\n",
       "  'attention_use_n_prev_rewards': 0,\n",
       "  'framestack': True,\n",
       "  'dim': 84,\n",
       "  'grayscale': False,\n",
       "  'zero_mean': True,\n",
       "  'custom_model': None,\n",
       "  'custom_model_config': {},\n",
       "  'custom_action_dist': None,\n",
       "  'custom_preprocessor': None,\n",
       "  'lstm_use_prev_action_reward': -1},\n",
       " 'optimizer': {},\n",
       " 'max_requests_in_flight_per_sampler_worker': 2,\n",
       " 'explore': True,\n",
       " 'exploration_config': {'type': 'StochasticSampling'},\n",
       " 'input_config': {},\n",
       " 'actions_in_input_normalized': False,\n",
       " 'postprocess_inputs': False,\n",
       " 'shuffle_buffer_size': 0,\n",
       " 'output': None,\n",
       " 'output_config': {},\n",
       " 'output_compress_columns': ['obs', 'new_obs'],\n",
       " 'output_max_file_size': 67108864,\n",
       " 'offline_sampling': False,\n",
       " 'evaluation_interval': None,\n",
       " 'evaluation_duration': 10,\n",
       " 'evaluation_duration_unit': 'episodes',\n",
       " 'evaluation_sample_timeout_s': 180.0,\n",
       " 'evaluation_parallel_to_training': False,\n",
       " 'evaluation_config': None,\n",
       " 'off_policy_estimation_methods': {},\n",
       " 'ope_split_batch_by_episode': True,\n",
       " 'evaluation_num_workers': 0,\n",
       " 'always_attach_evaluation_results': False,\n",
       " 'enable_async_evaluation': False,\n",
       " 'in_evaluation': False,\n",
       " 'sync_filters_on_rollout_workers_timeout_s': 60.0,\n",
       " 'keep_per_episode_custom_metrics': False,\n",
       " 'metrics_episode_collection_timeout_s': 60.0,\n",
       " 'metrics_num_episodes_for_smoothing': 100,\n",
       " 'min_time_s_per_iteration': None,\n",
       " 'min_train_timesteps_per_iteration': 0,\n",
       " 'min_sample_timesteps_per_iteration': 0,\n",
       " 'export_native_model_files': False,\n",
       " 'logger_creator': None,\n",
       " 'logger_config': None,\n",
       " 'log_level': 'WARN',\n",
       " 'log_sys_usage': True,\n",
       " 'fake_sampler': False,\n",
       " 'seed': None,\n",
       " 'worker_cls': None,\n",
       " '_tf_policy_handles_more_than_one_loss': False,\n",
       " '_disable_preprocessor_api': False,\n",
       " '_disable_action_flattening': False,\n",
       " '_disable_execution_plan_api': True,\n",
       " 'simple_optimizer': -1,\n",
       " 'replay_sequence_length': None,\n",
       " 'lr_schedule': None,\n",
       " 'use_critic': True,\n",
       " 'use_gae': True,\n",
       " 'kl_coeff': 0.2,\n",
       " 'sgd_minibatch_size': 128,\n",
       " 'num_sgd_iter': 30,\n",
       " 'shuffle_sequences': True,\n",
       " 'vf_loss_coeff': 1.0,\n",
       " 'entropy_coeff': 0.0,\n",
       " 'entropy_coeff_schedule': None,\n",
       " 'clip_param': 0.3,\n",
       " 'vf_clip_param': 10.0,\n",
       " 'grad_clip': None,\n",
       " 'kl_target': 0.01,\n",
       " 'vf_share_layers': -1,\n",
       " 'lambda': 1.0,\n",
       " 'input': 'sampler',\n",
       " 'multiagent': {'policies': {'default_policy': <ray.rllib.policy.policy.PolicySpec at 0x7f246ac90880>},\n",
       "  'policy_mapping_fn': <function ray.rllib.algorithms.algorithm_config.AlgorithmConfig.__init__.<locals>.<lambda>(aid, episode, worker, **kwargs)>,\n",
       "  'policies_to_train': None,\n",
       "  'policy_map_capacity': 100,\n",
       "  'policy_map_cache': None,\n",
       "  'count_steps_by': 'env_steps',\n",
       "  'observation_fn': None},\n",
       " 'callbacks': ray.rllib.algorithms.callbacks.DefaultCallbacks,\n",
       " 'create_env_on_driver': False,\n",
       " 'custom_eval_function': None,\n",
       " 'framework': 'tf',\n",
       " 'num_cpus_for_driver': 1,\n",
       " 'num_workers': 2}"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from ray.rllib.algorithms.ppo import PPOConfig\n",
    "from ray.rllib.algorithms.a2c import A2CConfig\n",
    "from ray.rllib.algorithms.td3 import TD3Config \n",
    "from ray.rllib.algorithms.dqn import DQNConfig \n",
    "\n",
    "PPOConfig().to_dict()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ONLY TRAINING WITHOUT TUNING"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* To just train your model, pick hard-coded hyperparameters\n",
    "\n",
    "### ALSO, INCREASE THE NUMBER OF TRAINING ITERATIONS AND NUM SAMPLES = 1, AS YOU HAVE ONE SET OF HYPERPARAMETERS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_ppo_params():\n",
    "    return {\n",
    "      \"entropy_coeff\": 0.0000001,\n",
    "      \"lr\": 5e-5,\n",
    "      \"sgd_minibatch_size\": 64,\n",
    "      \"lambda\":0.9,\"framework\":'torch'}\n",
    "\n",
    "drl_agent = DRLlibv2(\n",
    "    trainable=model_name,\n",
    "    train_env=env(train_env_config),\n",
    "    train_env_name=\"StockTrading_train\",\n",
    "    framework=\"torch\",\n",
    "    num_workers=1,\n",
    "    log_level=\"WARN\",\n",
    "    run_name = 'FINRL_TEST',\n",
    "    local_dir = \"FINRL_TEST\",\n",
    "    params = sample_ppo_params(),\n",
    "    num_samples = 1,\n",
    "    num_gpus=1,\n",
    "    training_iterations=5,\n",
    "    checkpoint_freq=5\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TUNING WITH MLP"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* In the model key, you have fcnet, which signifies MLP model.\n",
    "\n",
    "### ALSO, INCREASE THE NUMBER OF TRAINING ITERATIONS AND NUM SAMPLES\n",
    "\n",
    "* Training iterations will increase the number of training loops, and num samples is the number of hyperparameter configuration to try\n",
    "* Search algorithm picks the hyperparameter configuration based on Bayesian methods or Random methods. You can find more info [here](https://docs.ray.io/en/latest/tune/api/suggestion.html#tune-search-alg)\n",
    "\n",
    "* Scheduler stops inferior trials. You can find more info [here](https://docs.ray.io/en/latest/tune/api/schedulers.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def sample_ppo_params():\n",
    "  return {\n",
    "      \"entropy_coeff\": tune.loguniform(0.00000001, 0.1),\n",
    "      \"lr\": tune.loguniform(5e-5, 0.001),\n",
    "      \"sgd_minibatch_size\": tune.choice([ 32, 64, 128, 256, 512]),\n",
    "      \"lambda\": tune.choice([0.1,0.3,0.5,0.7,0.9,1.0]),\n",
    "      \"framework\":'torch',\n",
    "      'model':{\n",
    "        'fcnet_hiddens': [256, 256]\n",
    "      }\n",
    "  }\n",
    "\n",
    "search_alg = OptunaSearch(\n",
    "        metric=\"episode_reward_mean\",\n",
    "    mode=\"max\")\n",
    "scheduler_ = ASHAScheduler(\n",
    "        metric=\"episode_reward_mean\",\n",
    "        mode=\"max\",\n",
    "        max_t=5,\n",
    "        grace_period=1,\n",
    "        reduction_factor=2,\n",
    "    )\n",
    "\n",
    "drl_agent = DRLlibv2(\n",
    "    trainable=model_name,\n",
    "    train_env=env(train_env_config),\n",
    "    train_env_name=\"StockTrading_train\",\n",
    "    framework=\"torch\",\n",
    "    num_workers=1,\n",
    "    log_level=\"WARN\",\n",
    "    run_name = 'FINRL_TEST',\n",
    "    local_dir = \"FINRL_TEST\",\n",
    "    params = sample_ppo_params(),\n",
    "    num_samples = 1,\n",
    "    num_gpus=1,\n",
    "    training_iterations=5,\n",
    "    checkpoint_freq=5,\n",
    "    scheduler=scheduler_,\n",
    "    search_alg=search_alg\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TUNING WITH LSTM\n",
    "\n",
    "* To tune with LSTM, just pass `use_lstm` = True, and tune the lstm_cell_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_ppo_params():\n",
    "  return {\n",
    "      \"entropy_coeff\": tune.loguniform(0.00000001, 0.1),\n",
    "      \"lr\": tune.loguniform(5e-5, 0.001),\n",
    "      \"sgd_minibatch_size\": tune.choice([ 32, 64, 128, 256, 512]),\n",
    "      \"lambda\": tune.choice([0.1,0.3,0.5,0.7,0.9,1.0]),\n",
    "      \"framework\":'torch',\n",
    "      'model':{\n",
    "        'use_lstm':True,\n",
    "        'lstm_cell_size':tune.choice([128,256,512])\n",
    "      },\n",
    "  }\n",
    "\n",
    "search_alg = OptunaSearch(\n",
    "        metric=\"episode_reward_mean\",\n",
    "    mode=\"max\")\n",
    "scheduler_ = ASHAScheduler(\n",
    "        metric=\"episode_reward_mean\",\n",
    "        mode=\"max\",\n",
    "        max_t=5,\n",
    "        grace_period=1,\n",
    "        reduction_factor=2,\n",
    "    )\n",
    "\n",
    "drl_agent = DRLlibv2(\n",
    "    trainable=model_name,\n",
    "    train_env=env(train_env_config),\n",
    "    train_env_name=\"StockTrading_train\",\n",
    "    framework=\"torch\",\n",
    "    num_workers=1,\n",
    "    log_level=\"WARN\",\n",
    "    run_name = 'FINRL_TEST',\n",
    "    local_dir = \"FINRL_TEST\",\n",
    "    params = sample_ppo_params(),\n",
    "    num_samples = 1,\n",
    "    num_gpus=1,\n",
    "    training_iterations=5,\n",
    "    checkpoint_freq=5,\n",
    "    scheduler=scheduler_,\n",
    "    search_alg=search_alg\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TUNING WITH TRANSFORMER ARCHITECTURE (GTrXL)\n",
    "\n",
    "* To tune with transformers, just pass `use_attention=True`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_ppo_params():\n",
    "  return {\n",
    "      \"entropy_coeff\": tune.loguniform(0.00000001, 0.1),\n",
    "      \"lr\": tune.loguniform(5e-5, 0.001),\n",
    "      \"sgd_minibatch_size\": tune.choice([ 32, 64, 128, 256, 512]),\n",
    "      \"lambda\": tune.choice([0.1,0.3,0.5,0.7,0.9,1.0]),\n",
    "      'model':{\n",
    "        'use_attention': True,\n",
    "        'attention_num_transformer_units': 1,\n",
    "        'attention_dim': 64,\n",
    "        'attention_num_heads': 1,\n",
    "        'attention_head_dim': 32,\n",
    "        'attention_memory_inference': 50,\n",
    "        'attention_memory_training': 50,\n",
    "        'attention_position_wise_mlp_dim': 32,\n",
    "        'attention_init_gru_gate_bias': 2.0,\n",
    "        'attention_use_n_prev_actions': 0,\n",
    "        'attention_use_n_prev_rewards': 0,\n",
    "      }\n",
    "  }\n",
    "\n",
    "search_alg = OptunaSearch(\n",
    "        metric=\"episode_reward_mean\",\n",
    "    mode=\"max\")\n",
    "scheduler_ = ASHAScheduler(\n",
    "        metric=\"episode_reward_mean\",\n",
    "        mode=\"max\",\n",
    "        max_t=5,\n",
    "        grace_period=1,\n",
    "        reduction_factor=2,\n",
    "    )\n",
    "\n",
    "drl_agent = DRLlibv2(\n",
    "    trainable=model_name,\n",
    "    train_env=env(train_env_config),\n",
    "    train_env_name=\"StockTrading_train\",\n",
    "    framework=\"torch\",\n",
    "    num_workers=1,\n",
    "    log_level=\"WARN\",\n",
    "    run_name = 'FINRL_TEST',\n",
    "    local_dir = \"FINRL_TEST\",\n",
    "    params = sample_ppo_params(),\n",
    "    num_samples = 1,\n",
    "    num_gpus=1,\n",
    "    training_iterations=5,\n",
    "    checkpoint_freq=5,\n",
    "    scheduler=scheduler_,\n",
    "    search_alg=search_alg\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TRAIN AND TUNE THE MODEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"tuneStatus\">\n",
       "  <div style=\"display: flex;flex-direction: row\">\n",
       "    <div style=\"display: flex;flex-direction: column;\">\n",
       "      <h3>Tune Status</h3>\n",
       "      <table>\n",
       "<tbody>\n",
       "<tr><td>Current time:</td><td>2023-02-25 21:18:08</td></tr>\n",
       "<tr><td>Running for: </td><td>00:02:03.10        </td></tr>\n",
       "<tr><td>Memory:      </td><td>4.7/7.4 GiB        </td></tr>\n",
       "</tbody>\n",
       "</table>\n",
       "    </div>\n",
       "    <div class=\"vDivider\"></div>\n",
       "    <div class=\"systemInfo\">\n",
       "      <h3>System Info</h3>\n",
       "      Using FIFO scheduling algorithm.<br>Resources requested: 0/2 CPUs, 0/1 GPUs, 0.0/3.64 GiB heap, 0.0/1.82 GiB objects (0.0/1.0 accelerator_type:G)\n",
       "    </div>\n",
       "    \n",
       "  </div>\n",
       "  <div class=\"hDivider\"></div>\n",
       "  <div class=\"trialStatus\">\n",
       "    <h3>Trial Status</h3>\n",
       "    <table>\n",
       "<thead>\n",
       "<tr><th>Trial name                        </th><th>status    </th><th>loc             </th><th style=\"text-align: right;\">  iter</th><th style=\"text-align: right;\">  total time (s)</th><th style=\"text-align: right;\">   ts</th><th style=\"text-align: right;\">  reward</th><th style=\"text-align: right;\">  episode_reward_max</th><th style=\"text-align: right;\">  episode_reward_min</th><th style=\"text-align: right;\">  episode_len_mean</th></tr>\n",
       "</thead>\n",
       "<tbody>\n",
       "<tr><td>PPO_StockTrading_train_85c44_00000</td><td>TERMINATED</td><td>172.21.182.6:826</td><td style=\"text-align: right;\">     5</td><td style=\"text-align: right;\">         112.754</td><td style=\"text-align: right;\">20000</td><td style=\"text-align: right;\"> 977.338</td><td style=\"text-align: right;\">              1144.1</td><td style=\"text-align: right;\">              783.01</td><td style=\"text-align: right;\">              1907</td></tr>\n",
       "</tbody>\n",
       "</table>\n",
       "  </div>\n",
       "</div>\n",
       "<style>\n",
       ".tuneStatus {\n",
       "  color: var(--jp-ui-font-color1);\n",
       "}\n",
       ".tuneStatus .systemInfo {\n",
       "  display: flex;\n",
       "  flex-direction: column;\n",
       "}\n",
       ".tuneStatus td {\n",
       "  white-space: nowrap;\n",
       "}\n",
       ".tuneStatus .trialStatus {\n",
       "  display: flex;\n",
       "  flex-direction: column;\n",
       "}\n",
       ".tuneStatus h3 {\n",
       "  font-weight: bold;\n",
       "}\n",
       ".tuneStatus .hDivider {\n",
       "  border-bottom-width: var(--jp-border-width);\n",
       "  border-bottom-color: var(--jp-border-color0);\n",
       "  border-bottom-style: solid;\n",
       "}\n",
       ".tuneStatus .vDivider {\n",
       "  border-left-width: var(--jp-border-width);\n",
       "  border-left-color: var(--jp-border-color0);\n",
       "  border-left-style: solid;\n",
       "  margin: 0.5em 1em 0.5em 1em;\n",
       "}\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-02-25 21:18:08,638\tINFO tune.py:762 -- Total run time: 123.82 seconds (123.09 seconds for the tuning loop).\n"
     ]
    }
   ],
   "source": [
    "res = drl_agent.train_tune_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_df, best_result = drl_agent.infer_results()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GENERATE TEST ENVIRONMENT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "Shape of DataFrame:  (11340, 9)\n",
      "Clean data for AAPL\n",
      "Data clean for AAPL is finished.\n",
      "Clean data for AMGN\n",
      "Data clean for AMGN is finished.\n",
      "Clean data for AXP\n",
      "Data clean for AXP is finished.\n",
      "Clean data for BA\n",
      "Data clean for BA is finished.\n",
      "Clean data for CAT\n",
      "Data clean for CAT is finished.\n",
      "Clean data for CRM\n",
      "Data clean for CRM is finished.\n",
      "Clean data for CSCO\n",
      "Data clean for CSCO is finished.\n",
      "Clean data for CVX\n",
      "Data clean for CVX is finished.\n",
      "Clean data for DIS\n",
      "Data clean for DIS is finished.\n",
      "Clean data for DOW\n",
      "Data clean for DOW is finished.\n",
      "Clean data for GS\n",
      "Data clean for GS is finished.\n",
      "Clean data for HD\n",
      "Data clean for HD is finished.\n",
      "Clean data for HON\n",
      "Data clean for HON is finished.\n",
      "Clean data for IBM\n",
      "Data clean for IBM is finished.\n",
      "Clean data for INTC\n",
      "Data clean for INTC is finished.\n",
      "Clean data for JNJ\n",
      "Data clean for JNJ is finished.\n",
      "Clean data for JPM\n",
      "Data clean for JPM is finished.\n",
      "Clean data for KO\n",
      "Data clean for KO is finished.\n",
      "Clean data for MCD\n",
      "Data clean for MCD is finished.\n",
      "Clean data for MMM\n",
      "Data clean for MMM is finished.\n",
      "Clean data for MRK\n",
      "Data clean for MRK is finished.\n",
      "Clean data for MSFT\n",
      "Data clean for MSFT is finished.\n",
      "Clean data for NKE\n",
      "Data clean for NKE is finished.\n",
      "Clean data for PG\n",
      "Data clean for PG is finished.\n",
      "Clean data for TRV\n",
      "Data clean for TRV is finished.\n",
      "Clean data for UNH\n",
      "Data clean for UNH is finished.\n",
      "Clean data for V\n",
      "Data clean for V is finished.\n",
      "Clean data for VZ\n",
      "Data clean for VZ is finished.\n",
      "Clean data for WBA\n",
      "Data clean for WBA is finished.\n",
      "Clean data for WMT\n",
      "Data clean for WMT is finished.\n",
      "Data clean all finished!\n",
      "[*********************100%***********************]  1 of 1 completed\n",
      "Shape of DataFrame:  (378, 9)\n",
      "Clean data for ^VIX\n",
      "Data clean for ^VIX is finished.\n",
      "Data clean all finished!\n",
      "['AAPL' 'AMGN' 'AXP' 'BA' 'CAT' 'CRM' 'CSCO' 'CVX' 'DIS' 'DOW' 'GS' 'HD'\n",
      " 'HON' 'IBM' 'INTC' 'JNJ' 'JPM' 'KO' 'MCD' 'MMM' 'MRK' 'MSFT' 'NKE' 'PG'\n",
      " 'TRV' 'UNH' 'V' 'VZ' 'WBA' 'WMT']\n",
      "Successfully transformed into array\n"
     ]
    }
   ],
   "source": [
    "test_env_config = get_test_config(TEST_START_DATE, TEST_END_DATE, ticker_list, data_source, time_interval, \n",
    "                        technical_indicator_list, env, model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_env_instance = env(test_env_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-02-25 21:18:30,743\tWARNING algorithm_config.py:488 -- Cannot create PPOConfig from given `config_dict`! Property __stdout_file__ not supported.\n",
      "2023-02-25 21:18:30,757\tINFO algorithm.py:501 -- Current log_level is WARN. For more information, set 'log_level': 'INFO' / 'DEBUG' or use the -v and -vv flags.\n",
      "\u001b[2m\u001b[36m(RolloutWorker pid=1280)\u001b[0m 2023-02-25 21:18:32,889\tWARNING env.py:147 -- Your env doesn't have a .spec.max_episode_steps attribute. This is fine if you have set 'horizon' in your config dictionary, or `soft_horizon`. However, if you haven't, 'horizon' will default to infinity, and your environment will not be reset.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "test_agent = drl_agent.get_test_agent(test_env_instance,'StockTrading_testenv')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TESTING FOR MLP BASED AGENT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = test_env_instance.reset()\n",
    "\n",
    "episode_returns = list()  # the cumulative_return / initial_account\n",
    "episode_total_assets = list()\n",
    "episode_total_assets.append(test_env_instance.initial_total_asset)\n",
    "done = False\n",
    "while not done:\n",
    "    action  = test_agent.compute_single_action(observation=obs)\n",
    "    obs, reward, done, _ = test_env_instance.step(action)\n",
    "    # print(action)\n",
    "    total_asset = (\n",
    "        test_env_instance.amount\n",
    "        + (test_env_instance.price_ary[test_env_instance.day] * test_env_instance.stocks).sum()\n",
    "    )\n",
    "    episode_total_assets.append(total_asset)\n",
    "    episode_return = total_asset / test_env_instance.initial_total_asset\n",
    "    episode_returns.append(episode_return)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TESTING FOR LSTM BASED AGENT\n",
    "\n",
    "* THE ONLY THING THAT YOU HAVE TO TAKE CARE IS THE STATE SPACE SIZE\n",
    "\n",
    "```python\n",
    "    init_state = state = [\n",
    "    np.zeros([lstm_cell_size], np.float32) for _ in range(2)]\n",
    "```\n",
    "\n",
    "The testing agent expect the shape of the state as the above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = test_env_instance.reset()\n",
    "lstm_cell_size = res.get_best_result().config['model']['lstm_cell_size']\n",
    "init_state = state = [\n",
    "    np.zeros([lstm_cell_size], np.float32) for _ in range(2)]\n",
    "episode_returns = list()  # the cumulative_return / initial_account\n",
    "episode_total_assets = list()\n",
    "episode_total_assets.append(test_env_instance.initial_total_asset)\n",
    "done = False\n",
    "while not done:\n",
    "    action, state, _  = test_agent.compute_single_action(observation=obs,state=state)\n",
    "    obs, reward, done, _ = test_env_instance.step(action)\n",
    "    # print(action)\n",
    "    total_asset = (\n",
    "        test_env_instance.amount\n",
    "        + (test_env_instance.price_ary[test_env_instance.day] * test_env_instance.stocks).sum()\n",
    "    )\n",
    "    episode_total_assets.append(total_asset)\n",
    "    episode_return = total_asset / test_env_instance.initial_total_asset\n",
    "    episode_returns.append(episode_return)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TESTING FOR TRANSFORMER BASED AGENTS\n",
    "\n",
    "* HERE ALSO, YOU HAVE THE STATE SPACE FOR THE TRANSFORMER AS DIFFERENT\n",
    "\n",
    "```python\n",
    "    init_state = state = [\n",
    "     np.zeros([100, 64], np.float32) for _ in range(num_transformers) ]\n",
    "```\n",
    "\n",
    "Also, the state output is designed as the following\n",
    "\n",
    "```python\n",
    "    state = [\n",
    "        np.concatenate([state[i], [state_out[i]]], axis=0)[1:]\n",
    "        for i in range(num_transformers)\n",
    "    ]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = test_env_instance.reset()\n",
    "num_transformers = res.get_best_result().config[\"model\"][\"attention_num_transformer_units\"]\n",
    "\n",
    "init_state = state = [\n",
    "     np.zeros([100, 64], np.float32) for _ in range(num_transformers) ]\n",
    "episode_returns = list()  # the cumulative_return / initial_account\n",
    "episode_total_assets = list()\n",
    "episode_total_assets.append(test_env_instance.initial_total_asset)\n",
    "done = False\n",
    "while not done:\n",
    "    action, state_out, _  = test_agent.compute_single_action(observation=obs,state=state)\n",
    "    obs, reward, done, _ = test_env_instance.step(action)\n",
    "    # print(action)\n",
    "    total_asset = (\n",
    "        test_env_instance.amount\n",
    "        + (test_env_instance.price_ary[test_env_instance.day] * test_env_instance.stocks).sum()\n",
    "    )\n",
    "    episode_total_assets.append(total_asset)\n",
    "    episode_return = total_asset / test_env_instance.initial_total_asset\n",
    "    episode_returns.append(episode_return)\n",
    "    state = [\n",
    "        np.concatenate([state[i], [state_out[i]]], axis=0)[1:]\n",
    "        for i in range(num_transformers)\n",
    "    ]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## USING WEIGHTS AND BIASES FOR MLOPs ##TO-DO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mathe_kunal\u001b[0m (\u001b[33mrf_project\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import wandb\n",
    "\n",
    "wandb.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.air.integrations.wandb import WandbLoggerCallback\n",
    "\n",
    "wandb_callback = WandbLoggerCallback(project=\"Ray Tune Trial Run\",log_config=True,save_checkpoints=True)\n",
    "\n",
    "def sample_ppo_params():\n",
    "  return {\n",
    "      \"entropy_coeff\": tune.loguniform(0.00000001, 0.1),\n",
    "      \"lr\": tune.loguniform(5e-5, 0.001),\n",
    "      \"sgd_minibatch_size\": tune.choice([ 32, 64, 128, 256, 512]),\n",
    "      \"lambda\": tune.choice([0.1,0.3,0.5,0.7,0.9,1.0]),\n",
    "      'model':{\n",
    "        'use_attention': True,\n",
    "        'attention_num_transformer_units': 1,\n",
    "        'attention_dim': 64,\n",
    "        'attention_num_heads': 1,\n",
    "        'attention_head_dim': 32,\n",
    "        'attention_memory_inference': 50,\n",
    "        'attention_memory_training': 50,\n",
    "        'attention_position_wise_mlp_dim': 32,\n",
    "        'attention_init_gru_gate_bias': 2.0,\n",
    "        'attention_use_n_prev_actions': 0,\n",
    "        'attention_use_n_prev_rewards': 0,\n",
    "      }\n",
    "  }\n",
    "\n",
    "search_alg = OptunaSearch(\n",
    "        metric=\"episode_reward_mean\",\n",
    "    mode=\"max\")\n",
    "scheduler_ = ASHAScheduler(\n",
    "        metric=\"episode_reward_mean\",\n",
    "        mode=\"max\",\n",
    "        max_t=5,\n",
    "        grace_period=1,\n",
    "        reduction_factor=2,\n",
    "    )\n",
    "\n",
    "test_env_instance = env(test_env_config)\n",
    "\n",
    "drl_agent = DRLlibv2(\n",
    "    trainable=model_name,\n",
    "    train_env=env(train_env_config),\n",
    "    train_env_name=\"StockTrading_train\",\n",
    "    framework=\"torch\",\n",
    "    num_workers=1,\n",
    "    log_level=\"DEBUG\",\n",
    "    run_name = 'FINRL_TEST',\n",
    "    local_dir = \"FINRL_TEST\",\n",
    "    params = sample_ppo_params(),\n",
    "    num_samples = 1,\n",
    "    num_gpus=1,\n",
    "    training_iterations=5,\n",
    "    checkpoint_freq=5,\n",
    "    scheduler=scheduler_,\n",
    "    search_alg=search_alg,\n",
    "    callbacks=[wandb_callback]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mathe_kunal\u001b[0m (\u001b[33mrf_project\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.13.10"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/athekunal/Ray for FinRL/wandb/run-20230226_192559-vvjm51ve</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/rf_project/uncategorized/runs/vvjm51ve' target=\"_blank\">earnest-grass-6</a></strong> to <a href='https://wandb.ai/rf_project/uncategorized' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/rf_project/uncategorized' target=\"_blank\">https://wandb.ai/rf_project/uncategorized</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/rf_project/uncategorized/runs/vvjm51ve' target=\"_blank\">https://wandb.ai/rf_project/uncategorized/runs/vvjm51ve</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m:   6 of 6 files downloaded.  \n"
     ]
    }
   ],
   "source": [
    "import wandb\n",
    "run = wandb.init()\n",
    "artifact = run.use_artifact('rf_project/Ray Tune Trial Run/checkpoint_PPO_StockTrading_train_env_3e5ba38a:v0', type='model')\n",
    "artifact_dir = artifact.download()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from ray.rllib.algorithms import Algorithm\n",
    "\n",
    "# chkpt = \"/home/athekunal/Ray for FinRL/artifacts/checkpoint_PPO_StockTrading_train_env_3e5ba38a:v0/\"\n",
    "\n",
    "# # testing_agent = drl_agent.get_test_agent(test_env=test_env_instance,test_env_name=\"StockTest\",checkpoint=chkpt)\n",
    "\n",
    "# register_env(\"StockTest\", lambda: test_env_instance)\n",
    "# test_agent = Algorithm.from_checkpoint(checkpoint=chkpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
